
#include <cassert>
#include <cstdlib>

#include <DXGI.h>
#include <D3D10_1.h>
#include <D3DX11.h>
#include <D3DX10math.h>
#include <D3DX11async.h>

#include "Math\Matrix4x4.h"
#include "Math\Vector4.h"
#include "Scene\Material.h"
#include "Shaders\GBufferShader.h"
#include "Tools\RGBColor.h"

#pragma comment(lib, "D3D11.lib")
#pragma comment(lib, "D3DX11.lib")


ID3D10Blob* GBufferShader::s_vshader_buffer = NULL;
ID3D10Blob* GBufferShader::s_pshader_buffer = NULL;


BOOL GBufferShader::s_loaded = FALSE;


GBufferShader::GBufferShader()
	: m_vshader(NULL),
	  m_pshader(NULL),
      m_binded(FALSE),
	  m_frame_buffer(NULL),
	  m_material_buffer(NULL)
{
}


GBufferShader::~GBufferShader()
{	
	// Unbind if needed.
	if(m_binded)
		Unbind();
}


BOOL GBufferShader::IsLoaded()
{
	return s_loaded;
}


BOOL GBufferShader::Load()
{
	assert(!s_vshader_buffer);
	assert(!s_pshader_buffer);
	assert(!s_loaded);

	ID3D10Blob* error_buffer = NULL;
	LPCSTR compile_error = NULL;
	HRESULT result;

	// Load and compile pixel shader.
	result = D3DX11CompileFromFile(TEXT("Shaders\\Geometry.vs"),
								NULL, NULL, "Main", "vs_5_0",
								D3D10_SHADER_ENABLE_STRICTNESS, 0, NULL,
								&s_vshader_buffer, &error_buffer, NULL);

	if(FAILED(result))
		compile_error = static_cast< LPCSTR >(error_buffer->GetBufferPointer());

	assert(SUCCEEDED(result));

	// Load and compile pixel shader.
	result = D3DX11CompileFromFile(TEXT("Shaders\\GBuffer.ps"),
								NULL, NULL, "Main", "ps_5_0",
								D3D10_SHADER_ENABLE_STRICTNESS, 0, NULL,
								&s_pshader_buffer, &error_buffer, NULL);
	
	if(FAILED(result))
		compile_error = static_cast< LPCSTR >(error_buffer->GetBufferPointer());

	assert(SUCCEEDED(result));

	// Change state to loaded.
	s_loaded = TRUE;

	return TRUE;
}


VOID GBufferShader::Unload()
{
	assert(s_vshader_buffer);
	assert(s_pshader_buffer);
	assert(s_loaded);

	s_vshader_buffer->Release();
	s_vshader_buffer = NULL;

	s_pshader_buffer->Release();
	s_pshader_buffer = NULL;

	// Change state to unloaded.
	s_loaded = FALSE;
	s_pshader_buffer = NULL;
}


BOOL GBufferShader::IsBinded() CONST
{
	return m_binded;
}


BOOL GBufferShader::Bind()
{
	assert(s_vshader_buffer);
	assert(s_pshader_buffer);
	assert(s_loaded);
	assert(!m_vshader);
	assert(!m_pshader);
	assert(!m_binded);

	HRESULT result;

	// Create the vertex shader buffer.
	result = GetDevice()->CreateVertexShader(s_vshader_buffer->GetBufferPointer(),
											 s_vshader_buffer->GetBufferSize(),
											 NULL,
											 &m_vshader);

	assert(SUCCEEDED(result));

	// Create the pixel shader buffer.
	result = GetDevice()->CreatePixelShader(s_pshader_buffer->GetBufferPointer(),
											s_pshader_buffer->GetBufferSize(),
											NULL,
											&m_pshader);

	assert(SUCCEEDED(result));

	// Setup the input layout description.
	const UINT input_count = 2;
	D3D11_INPUT_ELEMENT_DESC input_desc[input_count];

	input_desc[0].SemanticName = "POSITION";
	input_desc[0].SemanticIndex = 0;
	input_desc[0].Format = DXGI_FORMAT_R32G32B32A32_FLOAT;
	input_desc[0].InputSlot = 0;
	input_desc[0].AlignedByteOffset = 0;
	input_desc[0].InputSlotClass = D3D11_INPUT_PER_VERTEX_DATA;
	input_desc[0].InstanceDataStepRate = 0;

	input_desc[1].SemanticName = "NORMAL";
	input_desc[1].SemanticIndex = 0;
	input_desc[1].Format = DXGI_FORMAT_R32G32B32_FLOAT;
	input_desc[1].InputSlot = 0;
	input_desc[1].AlignedByteOffset = D3D11_APPEND_ALIGNED_ELEMENT;
	input_desc[1].InputSlotClass = D3D11_INPUT_PER_VERTEX_DATA;
	input_desc[1].InstanceDataStepRate = 0;

	// Create the input layout.
	result = GetDevice()->CreateInputLayout(input_desc,
										    input_count,
									        s_vshader_buffer->GetBufferPointer(),
										    s_vshader_buffer->GetBufferSize(),
									        &m_layout);

	assert(SUCCEEDED(result));

	// Setup buffer descriptions.
	D3D11_BUFFER_DESC dynamic_desc;

	dynamic_desc.Usage = D3D11_USAGE_DYNAMIC;
	dynamic_desc.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
	dynamic_desc.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;
	dynamic_desc.MiscFlags = 0;
	dynamic_desc.StructureByteStride = 0;

	dynamic_desc.ByteWidth = sizeof(FrameBuffer);

	result = GetDevice()->CreateBuffer(&dynamic_desc, NULL, &m_frame_buffer);
	assert(SUCCEEDED(result));

	dynamic_desc.ByteWidth = sizeof(MaterialBuffer);

	result = GetDevice()->CreateBuffer(&dynamic_desc, NULL, &m_material_buffer);
	assert(SUCCEEDED(result));

	// Change state to binded.
	m_binded = TRUE;

	return m_binded;
}


VOID GBufferShader::Unbind()
{
	assert(m_layout);
	assert(m_vshader);
	assert(m_pshader);
	assert(m_binded);
	assert(m_frame_buffer);
	assert(m_material_buffer);

	m_vshader->Release();
	m_vshader = NULL;

	m_pshader->Release();
	m_pshader = NULL;

	m_layout->Release();
	m_layout = NULL;

	m_frame_buffer->Release();
	m_frame_buffer = NULL;

	m_material_buffer->Release();
	m_material_buffer = NULL;
	
	// Change state to unbinded.
	m_binded = FALSE;
}


VOID GBufferShader::Enable()
{
	assert(m_layout);
	assert(m_vshader);
	assert(m_pshader);
	assert(m_binded);

	// Enabling.
	GetContext()->IASetInputLayout(m_layout);
	GetContext()->VSSetShader(m_vshader, NULL, 0);
	GetContext()->PSSetShader(m_pshader, NULL, 0);
}


BOOL GBufferShader::SetFrameBuffer(CONST Vector4& position,
								   CONST Matrix4x4& world,
						           CONST Matrix4x4& view,
						           CONST Matrix4x4& projection)
{
	assert(m_layout);
	assert(m_vshader);
	assert(m_pshader);
	assert(m_binded);

	// Map matrices.
	D3D11_MAPPED_SUBRESOURCE mapped;
	
	HRESULT result = GetContext()->Map(m_frame_buffer,
									   0,
								       D3D11_MAP_WRITE_DISCARD,
								       0,
								       &mapped);

	assert(SUCCEEDED(result));

	FrameBuffer* ptr = reinterpret_cast< FrameBuffer* >(mapped.pData);

	// Position in world coordinates.
	Vector4 world_position = world * position;

	// Position copy.
	ptr->position[0] = world_position.GetX();
	ptr->position[1] = world_position.GetY();
	ptr->position[2] = world_position.GetZ();
	ptr->position[3] = world_position.GetW();

	// Matrices copy.
	for(UINT i = 0; i < 16; ++i)
	{
		UINT row = i / 4;
		UINT col = i % 4;

		ptr->world[i] = world(row, col);
		ptr->view[i] = view(row, col);
		ptr->projection[i] = projection(row, col);
	}

	GetContext()->Unmap(m_frame_buffer, 0);

	// Set matrices constant buffer.
	GetContext()->VSSetConstantBuffers(0, 1, &m_frame_buffer);

	return TRUE;
}


BOOL GBufferShader::SetMaterialBuffer(const Material& material)
{
	assert(m_layout);
	assert(m_vshader);
	assert(m_pshader);
	assert(m_binded);

	// Map material.
	D3D11_MAPPED_SUBRESOURCE mapped;
	
	HRESULT result = GetContext()->Map(m_material_buffer,
								       0,
								       D3D11_MAP_WRITE_DISCARD,
								       0,
								       &mapped);

	assert(SUCCEEDED(result));

	MaterialBuffer* ptr = reinterpret_cast< MaterialBuffer* >(mapped.pData);

	RGBColor ambient = material.GetAmbient();

	ptr->ambient[0] = ambient.GetR();
	ptr->ambient[1] = ambient.GetG();
	ptr->ambient[2] = ambient.GetB();

	RGBColor diffuse = material.GetDiffuse();

	ptr->diffuse[0] = diffuse.GetR();
	ptr->diffuse[1] = diffuse.GetG();
	ptr->diffuse[2] = diffuse.GetB();

	RGBColor specular = material.GetSpecular();

	ptr->specular[0] = specular.GetR();
	ptr->specular[1] = specular.GetG();
	ptr->specular[2] = specular.GetB();

	ptr->opacity = material.GetOpacity();
	ptr->shininess = material.GetShininess();

	GetContext()->Unmap(m_material_buffer, 0);

	// Set material constant buffer.
	GetContext()->PSSetConstantBuffers(0, 1, &m_material_buffer);

	return TRUE;
}